/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: alg coll remote lookup.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#include "alg_coll_remote_lookup.h"
#include "log.h"
#include "dtype_common.h"
#include "adapter_hccp.h"

namespace hccl {
using namespace std;

namespace {
    constexpr u32 LOOKUP_INFO_DATA_COUNT = 1;
    constexpr u32 LOOKUP_OFFSET_TRANSFER_RATE = 2;
}

enum HcclEsLookupShmIdx {
    LOOKUP_SHM_KEY_IDX = 0,
    LOOKUP_SHM_VALUE_IDX = 1,
    LOOKUP_SHM_IDX_NUM = 2
};

AlgCollRemoteLookUp::AlgCollRemoteLookUp() : funcTable_({
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOK_UP,
        std::bind(&AlgCollRemoteLookUp::AlgRemoteLookup, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::GET_LOOK_UP_REQUEST,
        std::bind(&AlgCollRemoteLookUp::AlgGetLookupRequest, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::WAIT_LOOK_UP_DATA,
        std::bind(&AlgCollRemoteLookUp::AlgWaitLookupData, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::SET_LOOK_UP_RESPONSE,
        std::bind(&AlgCollRemoteLookUp::AlgSetLookupResponse, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::WAIT_SOME,
        std::bind(&AlgCollRemoteLookUp::AlgWaitSome, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::SERVICE_CANCEL,
        std::bind(&AlgCollRemoteLookUp::AlgServiceCancel, this, std::placeholders::_1)}
    })
{
}

HcclResult AlgCollRemoteLookUp::Process(HcclOpPtr opPtr)
{
    CHK_PTR_NULL(opPtr);
    auto op = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    CHK_PRT_RET(funcTable_.find(op->param_.opType) == funcTable_.end(),
        HCCL_ERROR("[Process]funcTable_ without opType[%u]", op->param_.opType), HCCL_E_PARA);
    CHK_RET(funcTable_[op->param_.opType](op));
    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteLookUp::AlgGetLookupRequest(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    u32 psId = 0;
    bool notFindPsId = true;
    for (auto& pair : opPtr->param_.getRequestParam.psMap) {
        if (pair.second == opPtr->param_.getRequestParam.rankId) {
            psId = pair.first;
            notFindPsId = false;
            break;
        }
    }
    CHK_PRT_RET(notFindPsId, HCCL_ERROR("[AlgGetLookupRequest]psId Is not found.rankId:[%u]",
        opPtr->param_.getRequestParam.rankId), HCCL_E_NOT_FOUND);
    u32 workerSize = opPtr->param_.getRequestParam.workerList.size();
    if (workerSize == 0) {
        HCCL_WARNING("[Alg][GetLookupRequest] worker size is zero");
        return HCCL_SUCCESS;
    }

    u32 recvRequestCount = opPtr->param_.getRequestParam.lookupRecvRequestCount;
    HCCL_INFO("AlgGetLookupRequest dest rank %u work index %u",
        opPtr->param_.getRequestParam.workerList[(psId + recvRequestCount) % workerSize],
        (psId + recvRequestCount) % workerSize);
    TransTask *transTaskPtr = new (std::nothrow) TransTask();
    CHK_PTR_NULL(transTaskPtr);
    transTaskPtr->count = LOOKUP_INFO_DATA_COUNT;

    if (opPtr->param_.waitAny) {
        HCCL_INFO("AlgGetLookupRequest Using WaitAny!");
        transTaskPtr->improbeRanks = opPtr->param_.getRequestParam.workerList;
    } else {
        u32 improbeWorkerRank = opPtr->param_.getRequestParam.workerList[(psId + recvRequestCount) % workerSize];
        transTaskPtr->improbeRanks = {improbeWorkerRank};
        HCCL_DEBUG("current workerId[%u], tag[%d]", improbeWorkerRank, opPtr->param_.tag);
        opPtr->param_.getRequestParam.lookupReqStatus->workerId = improbeWorkerRank;
    }

    transTaskPtr->recvMem = opPtr->param_.keys;
    transTaskPtr->dataType = HCCL_DATA_TYPE_UINT64;
    opPtr->param_.getRequestParam.reqStatus->actualSize = 0;
    transTaskPtr->handle = opPtr->param_.getRequestParam.reqStatus;
    transTaskPtr->outputs.push_back(opPtr->param_.getRequestParam.reqStatus);
    transTaskPtr->outputs.push_back(&keyOffsetToValueOffset_);
    transTaskPtr->outputs.push_back(&actualKeyCount_);
    transTaskPtr->outputs.push_back(&opPtr->param_.getRequestParam.workerRankId);

    transTaskPtr->outputs.push_back(&opPtr->param_.getRequestParam.responseBuffer);

    transTaskPtr->taskType = OpTaskType::RECV;
    OpTaskPtr taskPtr(transTaskPtr);
    opPtr->PushTask(std::move(taskPtr));

    if (!opPtr->param_.waitAny) {
        opPtr->param_.getRequestParam.workerRankId =
            opPtr->param_.getRequestParam.workerList[(psId + recvRequestCount) % workerSize];
    }

    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteLookUp::AlgSetLookupResponse(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    void *sendMem = nullptr;
    if (actualKeyCount_ == 0) {
        HCCL_INFO("[AlgCollRemoteLookUp][SetLookupResponse] actual key count is zero");
    } else {
        u32 embeddingDim = opPtr->param_.valueCount / actualKeyCount_;

        if (opPtr->param_.keyMaxNum > ES_PIPELINE_THRESHOLD) {
            keyOffsetToValueOffset_ /= LOOKUP_OFFSET_TRANSFER_RATE;
        }
        HCCL_DEBUG("AlgCollRemoteLookUp: key offset[%llu] embeddingDim[%u] valueCount[%llu] actual key count[%llu]",
            keyOffsetToValueOffset_, embeddingDim, opPtr->param_.valueCount, actualKeyCount_);

        map<u32, vector<DeviceMem>> &transferMems = opPtr->param_.getRequestParam.transferMems;
        u32 workerRankId = opPtr->param_.getRequestParam.workerRankId;

        if (transferMems.find(workerRankId) != transferMems.end() &&
            transferMems[workerRankId].size() == LOOKUP_SHM_IDX_NUM) {
            OpTaskPtr taskPtr(new (std::nothrow) MemTask());
            CHK_SMART_PTR_NULL(taskPtr);
            MemTask *memTaskPtr = static_cast<MemTask *>(taskPtr.get());
            if (opPtr->param_.valuesType >= HCCL_DATA_TYPE_RESERVED) {
                HCCL_ERROR("[Get][DataSizeByType]errNo[0x%016llx] data type[%s] out of range[%d, %d]",
                    HCCL_ERROR_CODE(HCCL_E_INTERNAL), GetDataTypeEnumStr(opPtr->param_.valuesType).c_str(),
                    HCCL_DATA_TYPE_INT8, HCCL_DATA_TYPE_RESERVED - 1);
                return HCCL_E_PARA;
            }

            u64 keyBytesNum = keyOffsetToValueOffset_ * embeddingDim * SIZE_TABLE[opPtr->param_.valuesType];

            memTaskPtr->dstMem = static_cast<u8 *>(transferMems[workerRankId][LOOKUP_SHM_VALUE_IDX].ptr()) +
                keyBytesNum;
            memTaskPtr->maxSize = transferMems[workerRankId][LOOKUP_SHM_VALUE_IDX].size() - keyBytesNum;

            memTaskPtr->srcMem = opPtr->param_.values;
            memTaskPtr->size = opPtr->param_.valueCount * SIZE_TABLE[opPtr->param_.valuesType];
            memTaskPtr->taskType = OpTaskType::MEMCPY;
            sendMem = memTaskPtr->dstMem;

            opPtr->PushTask(std::move(taskPtr));

            HCCL_DEBUG("TransferMems is not empty. workerRankId[%u]", workerRankId);
        }

        HCCL_DEBUG("WorkerRankId[%u], transferMems size[%u]", workerRankId, transferMems.size());
    }

    TransTask *transTaskPtr = new (std::nothrow) TransTask();
    CHK_PTR_NULL(transTaskPtr);
    transTaskPtr->count = opPtr->param_.valueCount;
    transTaskPtr->dataType = opPtr->param_.valuesType;
    transTaskPtr->sendMem = sendMem == nullptr ? opPtr->param_.values : sendMem;
    transTaskPtr->outRequest = opPtr->param_.getRequestParam.outRequest;
    transTaskPtr->taskType = OpTaskType::ISEND;
    transTaskPtr->tag = opPtr->param_.tag;

    transTaskPtr->responseBuffer = opPtr->param_.getRequestParam.responseBuffer;

    OpTaskPtr taskPtr(transTaskPtr);
    opPtr->PushTask(std::move(taskPtr));

    HCCL_INFO("Isend lookup response, dest rank %u",  opPtr->param_.getRequestParam.workerRankId);
    return HCCL_SUCCESS;
}
}
